"""UNet++推理脚本"""
import os, sys, argparse, torch
sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))

from baselines.unetpp.model import UNetPlusPlus
from baselines.unetpp.config import MODEL_CONFIG, DATA_CONFIG
from utils.infer_utils import inference_single_image, load_model_for_inference

def main():
    parser = argparse.ArgumentParser(description='UNet++推理脚本')
    parser.add_argument('--image', type=str, required=True)
    parser.add_argument('--checkpoint', type=str, required=True)
    parser.add_argument('--output_dir', type=str, default='outputs/unetpp')
    parser.add_argument('--threshold', type=float, default=0.5)
    parser.add_argument('--device', type=str, default='cuda')
    args = parser.parse_args()

    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    model = UNetPlusPlus(**MODEL_CONFIG)
    model = load_model_for_inference(model, args.checkpoint, device)

    image_name_no_ext = os.path.splitext(os.path.basename(args.image))[0]
    output_path = os.path.join(args.output_dir, f'{image_name_no_ext}_pred.png')

    inference_time, foreground_pixels = inference_single_image(
        model, args.image, output_path, device, DATA_CONFIG['to_rgb'], args.threshold
    )

    print(f'[INFO] 推理用时: {inference_time:.4f}秒, 前景像素数: {foreground_pixels}')

if __name__ == '__main__':
    main()
